"""PointTrasnformer Test"""

import os, sys, argparse
sys.path.append(os.path.abspath(os.path.join(os.getcwd())))
path = os.path.abspath(os.path.join(os.getcwd()))

from mind3d.models.blocks.TransitionDown import Transitiondown
from mind3d.models.blocks.Transformer import Transformer
from mind3d.models.blocks.TransitionUp import TransitionUp
import mindspore
import numpy as np

def test_Transformer():
    X = mindspore.Tensor(np.random.rand(32, 1024, 3), dtype = mindspore.float32)
    Y = mindspore.Tensor(np.random.rand(32, 1024, 32), dtype = mindspore.float32)
    modeule = Transformer(32,512,16)
    Output, attn = modeule(X, Y)
    assert Output.shape == (32, 1024,32), 'output shape not match'
    assert attn.shape == (32, 1024, 16, 512), 'attention shape not match'
    print("PointTransformer Transformer block test passed!")

def test_TransitionDown():
    X = mindspore.Tensor(np.random.rand(32, 1024, 3), dtype = mindspore.float32)
    Y = mindspore.Tensor(np.random.rand(32, 1024, 32), dtype = mindspore.float32)
    channel = 32 * 2 ** 1
    module = Transitiondown(1024 // 4 ** 1, 16,[channel // 2 + 3, channel, channel])
    xyz, points= module(X,Y)
    assert xyz.shape == (32, 256, 3), 'xyz shape not match'
    assert points.shape == (32, 256, 64), 'points shape not match'
    print("PointTransformer TransitionDown block test passed!")


def test_TransitionUp():
    XYZ1 = mindspore.Tensor(np.random.rand(32, 4, 3), dtype = mindspore.float32)
    XYZ2 = mindspore.Tensor(np.random.rand(32, 4, 512), dtype = mindspore.float32)
    Points1 = mindspore.Tensor(np.random.rand(32, 16, 3), dtype = mindspore.float32)
    Points2 = mindspore.Tensor(np.random.rand(32, 16, 256), dtype = mindspore.float32)
    module = TransitionUp(512,256,256)
    Output = module(XYZ1,XYZ2,Points1,Points2)
    assert Output.shape == (32, 16, 256), 'output shape not match'
    print("PointTransformer TransitionUp block test passed!")

if __name__ == "__main__":
    test_TransitionDown()
    